import requests
import trafilatura
import akshare as ak
from tavily import TavilyClient
from langchain_core.tools import Tool
from pydantic import BaseModel, Field
from src.common.logger import getLogger
from tencentcloud.common import credential
from tencentcloud.tmt.v20180321 import tmt_client, models

logger = getLogger()

class SearchWebSchema(BaseModel):
    query: str = Field(description = "要搜索的内容或关键词")

class CalculateSchema(BaseModel):
    expression: str = Field(description = "数学表达式，如'3 + 5 * 2'")

class SearchWeatherSchema(BaseModel):
    city: str = Field(description = "查询天气的城市名")

class StockPriceSchema(BaseModel):
    symbol: str = Field(description = "沪深300指数的股票代码")
    query_date: str = Field(description = "查询日期")

class ExtractWebsiteTextSchema(BaseModel):
    url: str = Field(description = "资源的网络地址")

class TranslateTextSchema(BaseModel):
    text: str = Field(description = "需要翻译的文本")
    src: str = Field(description = "文本的语言类别，如'zh'、'en'")
    dest: str = Field(description = "翻译后的语言类别，如'zh'、'en'")

class AgentTools:

    def __init__(self):
        self.travil_api_key = "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"
        self.tengxun_sectet_id = "AKIDkRwKUSDl7FAxXu3yQipohNq7EXnDXgif"
        self.tengxun_sectet_key = "pTwEMI3eCallbGhxDJCuRp1zmU0JSRFp"

    def calculate_numerical(self, expression):
        logger.info(f"AgentTools calculate_numerical expression: {expression}")
        try:
            result = str(eval(expression))
            logger.info(f"AgentTools calculate_numerical result: {result}")
            return result
        except Exception as e:
            logger.error(f"AgentTools calculate_numerical error: {e}")
            raise Exception("计算表达式无效")

    def search_weather(self, city):
        logger.info(f"AgentTools search_weather city: {city}")
        try:
            endpoint = "https://wttr.in"
            response = requests.get(f"{endpoint}/{city}")
            response.raise_for_status()
            text_result = response.text
            logger.info(f"AgentTools search_weather text_result: \n{text_result}")
            return text_result
        except Exception as e:
            logger.error(f"AgentTools search_weather error: {e}")
            raise Exception("无法获取给定城市的天气情况")

    def snatch_stock_price(self, symbol, query_date):
        logger.info(f"snatch_stock_price symbol: {symbol}, query_date: {query_date}")
        try:
            query_date = query_date.replace("-", "")
            df = ak.stock_zh_a_hist(symbol = symbol, period = "daily", start_date = query_date, end_date = query_date, adjust = "qfq")
            logger.info(f"snatch_stock_price df: \n{df}")
            if not df.empty:
                df.sort_values(by = '日期', ascending = False, inplace = True)
                output = f"沪深300指数的股票: {symbol} 在 {query_date} 的开盘价为: {df.iloc[0]["开盘"]}，收盘价为：{df.iloc[0]["收盘"]}"
                logger.info(output)
                return output
            else:
                return "未获取到数据，请检查日期或网络连接"
        except Exception as e:
            logger.error(f"snatch_stock_price Error: {e}")
            return f"获取数据失败: {e}"

    def translate_text(self, text, source: str = "zh", target: str = "en"):
        logger.info(f"AgentTools translate_text text len: {len(text)}, source: {source}, target: {target}")
        cred = credential.Credential(self.tengxun_sectet_id, self.tengxun_sectet_key)
        client = tmt_client.TmtClient(cred, "ap-beijing")
        req = models.TextTranslateRequest()
        req.SourceText = text
        req.Source = source
        req.Target = target
        req.ProjectId = 0
        resp = client.TextTranslate(req)
        result = resp.TargetText
        logger.info(f"AgentTools translate_text result：{result}")
        return resp.TargetText

    def extract_website_text(self, url):
        logger.info(f"AgentTools extract_website_text url: {url}")
        website_text = ""
        downloader = trafilatura.fetch_url(url)
        if downloader:
            content = trafilatura.extract(
                downloader,
                include_comments = False,  # 不包含评论
                include_tables = True,  # 可选：保留表格
                output_format = "txt",  # 纯文本（也可选 'xml', 'json'）
                favor_precision = True  # 更精准，可能少一点内容但更干净
            )
            if content:
                logger.info(f"AgentTools extract_website_text content len: {len(content)}")
                website_text += content.strip()
        logger.info(f"AgentTools extract_website_text website_text len: {len(website_text)}")
        return website_text

    def search_web(self, query):
        logger.info(f"AgentTools search_web query: {query}")
        max_results = 5
        travil_client = TavilyClient(api_key = self.travil_api_key)
        search_results = travil_client.search(query = query, search_depth = "advanced", include_domains = None, max_results = max_results)
        logger.info(f"AgentTools search_web search_results len: {len(search_results)}")

        web_text = ""
        for result in search_results["results"]:
            url = result["url"]
            title = result.get("title", "")
            logger.info(f"AgentTools search_web title: {title}, url: {url}")
            downloader = trafilatura.fetch_url(url)
            if downloader:
                content = trafilatura.extract(
                    downloader,
                    include_comments = False,  # 不包含评论
                    include_tables = True,  # 可选：保留表格
                    output_format = "txt",  # 纯文本（也可选 'xml', 'json'）
                    favor_precision = True  # 更精准，可能少一点内容但更干净
                )
                if content:
                    logger.info(f"AgentTools search_web content len: {len(content)}")
                    web_text += content.strip()
        logger.info(f"AgentTools search_web web_text len: {len(web_text)}")
        return web_text

    def get_execute_tools(self):
        calc_tool = Tool(
            name = "calculate_numerical",
            description = "用于执行数学计算，只能用于数值计算",
            func = self.calculate_numerical,
            args_schema = CalculateSchema
        )
        weather_tool = Tool(
            name = "search_weather",
            description = "用于查询城市当天及未来的实时天气情况，禁止用来查询过去时间的天气情况",
            func = self.search_weather,
            args_schema = SearchWeatherSchema
        )
        stock_tool = Tool(
            name = "snatch_stock_price",
            description = "用于查询沪深300指数指定股票的开盘/收盘价",
            func = self.snatch_stock_price,
            args_schema = StockPriceSchema
        )
        translate_tool = Tool(
            name = "translate_text",
            description = "翻译文本为指定的语言",
            func = self.translate_text,
            args_schema = TranslateTextSchema
        )
        extract_tool = Tool(
            name = "extract_website_text",
            description = "抓取网页地址的内容",
            func = self.extract_website_text,
            args_schema = ExtractWebsiteTextSchema
        )
        web_tool = Tool(
            name = "search_web",
            description = "用于搜索真实世界的信息，如事实、人物、地点等",
            func = self.search_web,
            args_schema = SearchWebSchema
        )
        return [ calc_tool, weather_tool, stock_tool, extract_tool, translate_tool, web_tool ]
